import argparse
from fastapi import FastAPI
from fastapi.responses import Response  # 导入 Response 类
import uvicorn
import edge_tts
from edge_tts import VoicesManager
from pydantic import BaseModel

app = FastAPI()

# 定义TTS请求接口参数
class RequestInfer(BaseModel):
    text: str = "No text requiring voice synthesis was found."                         # 需要合成的文本
    voice_name: str = ""                                                               # 要使用哪个声音来生成
    output_audio: str = ""                                                        # 输出mp3文件

# test
@app.get("/tts")
async def root():
    return {"message": "This is Edge-TTS inference API. If you see this page, it means the server is running."}

# 获取所有声音类型列表
@app.get("/tts/edge/getVoiceList")
async def getVoicesList():
    voices_manager = await VoicesManager.create()
    # voice = voices.find(Gender="Male", Language="es")
    print("return VoiceList.")
    return voices_manager.voices

# 获取所有声音类型列表
@app.get("/tts/edge/getVoiceListBy")
async def getVoicesListBy(
    gender: str = None,
    locale: str = None,
):
    voices_manager = await VoicesManager.create()
    return voices_manager.find(Gender=gender, Locale=locale)
    # return voices_manager.find(Locale="en-US", Gender="Female", Language="en")   # 注意Locale后面带后缀，Male/Female要注意大小写，查询条件是依据返回值中的对象的字段名。

# post请求，本地生成mp3视频
@app.post("/tts/edge/gen_mp3")
async def edgetts_speech(request: RequestInfer):
    try:
        communicate = edge_tts.Communicate(request.text, request.voice_name)
        await communicate.save(request.output_audio + ".mp3")
        print("success for gen mp3.")
        # await communicate.save("spanish.mp3")

    except Exception as e:
        print(e)
        return {
            "error": "parameter_error"
        }
    
# post请求，返回二进制流
@app.post("/tts/edge/speech")
async def edgetts_speech(request: RequestInfer) -> bytes:
    communicate = edge_tts.Communicate(request.text, request.voice_name)
    # 用于存储二进制语音数据的字节列表
    audio_bytes = bytearray()
    # 异步迭代获取语音片段
    async for chunk in communicate.stream():
        # 只处理包含音频数据的块（跳过元数据）
        if chunk["type"] == "audio":
            audio_bytes.extend(chunk["data"])
    print("TTS success ：" + request.text)
    # print(f"生成的语音数据大小: {len(bytes(audio_bytes))} 字节")
    
    # 可选：保存为本地文件验证结果
    # with open(request.output_audio, "wb") as f:
    #     f.write(bytes(audio_bytes))
    # print("语音已保存为 output.mp3")
    
    # 使用 Response 类包装二进制数据，并指定媒体类型为 MP3
    return Response(
        content=bytes(audio_bytes),
        media_type="audio/mpeg"  # MP3 格式的标准媒体类型
    )


def main() -> None:
    global port
    parser = argparse.ArgumentParser(description="TTS Inference API")
    parser.add_argument("-p","--port", type=int, default=9880, help="端口")
    args = parser.parse_args()
    port = args.port

    print(f"service start: http://127.0.0.1:{port}/tts")


    # webbrowser.open(f"http://127.0.0.1:{port}")
    uvicorn.run(app=app, host="0.0.0.0", port=port, log_level="error")

if __name__ == "__main__":
    main()